﻿//-----------------------------------------------------------------------------------------------------------
// WBFSSync Project by Omega Frost 
// http://wbfssync.codeplex.com/
//
// WBFSSync is Licensed under the terms of the 
// Microsoft Reciprocal License (Ms-RL)
//-----------------------------------------------------------------------------------------------------------
using System;
using System.IO;
using System.Net;
using System.Security.Cryptography;
using WBFSe3.IO;

namespace WBFSe3.Wbfs
{
    //-------------------------------------------------------------------------------------------------------
    //
    //-------------------------------------------------------------------------------------------------------
    public class WbfsBuilder
    {
        //---------------------------------------------------------------------------------------------------
        //
        //---------------------------------------------------------------------------------------------------
        public struct PartitionInfo
        {
            //Informações básicas
            public WbfsPartitionType Type;
            public int Index;
            public long Position;
        }


        //---------------------------------------------------------------------------------------------------
        //
        //---------------------------------------------------------------------------------------------------
        public struct PartitionData
        {
            //Informações básicas
            public PartitionInfo Info;
            public int DataStartBlock;

            //Dados de criptografia e arquivos
            public byte[] Ticket;
            public int TMDSize;
            public long TMDOffset;
            public int CertChainSize;
            public long CertChainOffset;
            public long H3TableOffset; //H3 Table Size = 0x18000
            public long DataOffset;
            public long DataSize;
            public byte[] TMD;
            public byte[] Cert;
            public byte[] TitleKey;

            public int DolOffset;
            public int FstOffset;
            public int FstSize;
            public int AplOffset;
            public int AplSize;

            public byte[] AplHeader;
            public byte[] Fst;

            public int FstItemsCount;
            public int FstNameTablePos;
        }


        //---------------------------------------------------------------------------------------------------
        //
        //---------------------------------------------------------------------------------------------------
        public struct ZoneInfo
        {
            public WbfsPartitionType Type;
            public int StartSector;
            public int SectorCount;
        }

        
        //---------------------------------------------------------------------------------------------------
        //
        //---------------------------------------------------------------------------------------------------

        //Leitor do disco e de criptografia
        public BinaryReader reader;
        public Rijndael aesDecryptor;

        //Buffer de leitura
        public byte[] ReadBlock;

        //Código de verificação
        public uint magic1;
        public uint magic2;

        //Informações do disco
        public String Name;
        public String Code;
        
        public byte[] NameData;
        public byte[] CodeData;
        public byte[] RegionData;

        public WbfsRegion CodeRegion;
        public WbfsRegion DiscRegion;

        //Flags da partição
        public bool HasUpdate;
        public bool HasChannel;

        //Uso de setores
        public int UsedWiiSectors;
        public Boolean[] WiiSectorUsage;
        public ZoneInfo[] Zones;

        //Informação das Partições
        public int PartitionCount;
        public long PartitionInfoTablePosition;
        public PartitionData CurrentPartition;
        public PartitionInfo[] Partitions;


        //---------------------------------------------------------------------------------------------------
        // 
        //---------------------------------------------------------------------------------------------------
        public WbfsBuilder(Stream stream)
        {
            //Cria o leitor, a tabela de uso de setores 
            //e o buffer de leitura
            this.reader = new BinaryReader(stream);
            this.ReadBlock = new byte[WbfsConstants.WiiSectorSize];
            this.WiiSectorUsage = new Boolean[WbfsConstants.WiiSectorsPerDisc];
        }


        //---------------------------------------------------------------------------------------------------
        // Builds the Wii Sisc
        //---------------------------------------------------------------------------------------------------
        public WbfsError Build(Boolean buildAll)
        {
            WbfsError e = WbfsError.Ok;

            try
            {
                // The header is always used
                this.WiiSectorUsage[0] = true;

                if ((e = ParseHeader()) != WbfsError.Ok) return e;
                if ((e = ParsePartitionTable()) != WbfsError.Ok) return e;
                if ((e = ParseRegion()) != WbfsError.Ok) return e;
                if ((e = ParseSecondMagic()) != WbfsError.Ok) return e;

                // Build the sector usage
                if (buildAll)
                {
                    //Cria uma mapa contendo o tipo dos setores na tabela
                    this.Zones = new ZoneInfo[this.PartitionCount + 1];

                    //Marca os setores essenciais do disco (cabeçalho, tabela de partições e código de verificação)
                    this.Zones[0].Type = WbfsPartitionType.Data;
                    this.Zones[0].StartSector = 0;
                    this.Zones[0].SectorCount = GetLastUsedWiiSector() - this.Zones[0].StartSector;

                    for (int i = 0; i < this.PartitionCount; i++)
                    {
                        // Skip the channel installer until i understand it...
                        if (this.Partitions[i].Type == WbfsPartitionType.ChannelInstaler)
                            continue;

                        // Prepare a new partition
                        this.CurrentPartition = new PartitionData();
                        this.CurrentPartition.Info = this.Partitions[i];

                        // Parse the patition
                        if ((e = ParsePartition()) != WbfsError.Ok)
                            return e;

                        // Mark the partition as used
                        this.Zones[i + 1].Type = this.CurrentPartition.Info.Type;
                        this.Zones[i + 1].StartSector = (int)(this.CurrentPartition.Info.Position / WbfsConstants.WiiSectorSize);
                        this.Zones[i + 1].SectorCount = GetLastUsedWiiSector() - this.Zones[i + 1].StartSector;
                    }

                    // Build the usage table
                    this.UsedWiiSectors = 0;
                    for (int i = 0; i < WbfsConstants.WiiSectorsPerDisc; i++)
                    {
                        if (this.WiiSectorUsage[i])
                            this.UsedWiiSectors++;
                    }
                }
            }
            catch (WbfsIOException exception)
            {
                return (WbfsError)exception.NativeError;
            }
            catch (Exception)
            {
                return WbfsError.Fail;
            }

            return WbfsError.Ok;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads the Disc Header
        //---------------------------------------------------------------------------------------------------
        private WbfsError ParseHeader()
        {
            this.CodeData = new byte[6];
            this.NameData = new byte[64];

            //Code
            this.reader.BaseStream.Position = WbfsConstants.WiiCodePos;
            
            if (this.reader.Read(this.CodeData, 0, WbfsConstants.WiiCodeLen) != WbfsConstants.WiiCodeLen)
                return WbfsError.IOReadError;

            this.CodeRegion = WbfsHelper.GetRegionFromCode(this.CodeData[3]);
            this.Code = WbfsHelper.ReadString(this.CodeData, 0);

            //Magic 1
            this.reader.BaseStream.Position = WbfsConstants.WiiMagic1Pos;
            this.magic1 = this.ReadUInt32();

            if (this.magic1 != WbfsConstants.WiiMagic1)
                return WbfsError.DiscMagic1Incorrect;

            //Name
            this.reader.BaseStream.Position = WbfsConstants.WiiNamePos;

            if (this.reader.Read(this.NameData, 0, WbfsConstants.WiiNameLen) != WbfsConstants.WiiNameLen)
                return WbfsError.IOReadError;

            this.Name = WbfsHelper.ReadString(this.NameData, 0);

            return WbfsError.Ok;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads the disc's partition table
        //---------------------------------------------------------------------------------------------------
        private WbfsError ParsePartitionTable()
        {
            //Marca a região da tabela de partição como usada
            //e avança a posição do stream
            this.WiiSectorUsage[8] = true; //0x40000 / 0x8000
            this.reader.BaseStream.Position = 0x40000;

            //Le o número de partições, a posição da tabela
            //e aloca o vetor de informações
            this.PartitionCount = this.ReadInt32();
            this.PartitionInfoTablePosition = (long)this.ReadInt32() << 2;
            this.Partitions = new PartitionInfo[this.PartitionCount];

            //Vai até a tabela de partições e marca como usada
            this.reader.BaseStream.Position = this.PartitionInfoTablePosition;
            this.WiiSectorUsage[(int)(this.PartitionInfoTablePosition / WbfsConstants.WiiSectorSize)] = true;

            //Lê a tabela de partições
            for (int i = 0; i < this.PartitionCount; i++)
            {
                this.Partitions[i].Position = (long)this.ReadInt32() << 2;
                this.Partitions[i].Type = (WbfsPartitionType)this.ReadInt32();

                switch (this.Partitions[i].Type)
                {
                    case WbfsPartitionType.Update:
                        this.HasUpdate = true;
                        break;
                    case WbfsPartitionType.ChannelInstaler:
                        this.HasChannel = true;
                        break;
                }
            }

            return WbfsError.Ok;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads the disc region codes
        //---------------------------------------------------------------------------------------------------
        private WbfsError ParseRegion()
        {
            //Marca o setor como usado
            this.WiiSectorUsage[9] = true; //0x4E000 / 0x8000

            this.RegionData = new byte[32];

            this.reader.BaseStream.Position = 0x4E000;
            this.reader.Read(this.RegionData, 0, 32);

            bool IsJap = (sbyte)this.RegionData[16] > 0;
            bool IsUsa = (sbyte)this.RegionData[17] > 0;
            bool IsEur = (sbyte)this.RegionData[20] > 0;

            if (IsUsa)
                this.DiscRegion = WbfsRegion.NTSCU;
            else if (IsEur)
                this.DiscRegion = WbfsRegion.PAL;
            else if (IsJap)
                this.DiscRegion = WbfsRegion.NTSCJ;
            else
            {
                int RegionByte = WbfsHelper.ntohi(this.RegionData, 0);
                this.DiscRegion = (WbfsRegion)RegionByte;
            }

            return WbfsError.Ok;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads the second verification code
        //---------------------------------------------------------------------------------------------------
        private WbfsError ParseSecondMagic()
        {
            //Ao contrário do Magic 1, esse não vai falhar o processo, apenas dará o aviso

            //Magic 2
            this.WiiSectorUsage[9] = true; //0x4FFFC / 0x8000
            this.reader.BaseStream.Position = 0x4FFFC;

            this.magic2 = this.ReadUInt32();
            if (this.magic2 != 0xC3F81A8E)
            {
                //???
            }

            return WbfsError.Ok;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads a Disc partition
        //---------------------------------------------------------------------------------------------------
        private WbfsError ParsePartition()
        {
            WbfsError e = WbfsError.Ok;

            //Ticket
            this.CurrentPartition.Ticket = new byte[0x2A4];
            ReadPartition(this.CurrentPartition.Ticket, 0, 0x2A4);

            //TMD
            this.CurrentPartition.TMDSize = this.ReadInt32();
            this.CurrentPartition.TMDOffset = (long)this.ReadInt32() << 2;

            //Cert
            this.CurrentPartition.CertChainSize = this.ReadInt32();
            this.CurrentPartition.CertChainOffset = (long)this.ReadInt32() << 2;

            //H3
            this.CurrentPartition.H3TableOffset = (long)this.ReadInt32() << 2;

            //Data
            this.CurrentPartition.DataOffset = (long)this.ReadInt32() << 2;
            this.CurrentPartition.DataSize = (long)this.ReadInt32() << 2;

            //Bloco
            this.CurrentPartition.DataStartBlock = (int)((this.CurrentPartition.Info.Position +
                this.CurrentPartition.DataOffset) >> WbfsConstants.WiiSectorSizeLSh);

            //TMD
            this.CurrentPartition.TMD = new byte[this.CurrentPartition.TMDSize];
            if ((e = ReadPartition(
                this.CurrentPartition.TMD,
                this.CurrentPartition.TMDOffset,
                this.CurrentPartition.TMDSize)) != WbfsError.Ok)
                goto RETURN;

            //Cert
            this.CurrentPartition.Cert = new byte[this.CurrentPartition.CertChainSize];
            if ((e = ReadPartition(
                this.CurrentPartition.Cert,
                this.CurrentPartition.CertChainOffset,
                this.CurrentPartition.CertChainSize)) != WbfsError.Ok)
                goto RETURN;

            //Title Key
            if ((e = ParseTitleKey()) != WbfsError.Ok) goto RETURN;

            //Marca a tabela H3 como região utilizada do disco
            MarkUsedSectors(
                this.CurrentPartition.Info.Position +
                this.CurrentPartition.H3TableOffset,
                0x18000);

            //Processa o conteúdo propriamente dito da partição (criptografado)
            e = ReadPartitionData(this.ReadBlock, 0, 0x7C00);

            this.CurrentPartition.DolOffset = WbfsHelper.ntohi(this.ReadBlock, 0x420) << 2;
            this.CurrentPartition.FstOffset = WbfsHelper.ntohi(this.ReadBlock, 0x424) << 2;
            this.CurrentPartition.FstSize = WbfsHelper.ntohi(this.ReadBlock, 0x428) << 2;
            this.CurrentPartition.AplOffset = 0x2440;

            this.CurrentPartition.AplHeader = new byte[0x20];
            ReadPartitionData(
                this.CurrentPartition.AplHeader,
                this.CurrentPartition.AplOffset,
                0x20);

            this.CurrentPartition.AplSize = 0x20 +
                WbfsHelper.ntohi(this.CurrentPartition.AplHeader, 0x14) +
                WbfsHelper.ntohi(this.CurrentPartition.AplHeader, 0x18);

            //Marca os blocos APL e DOL como regiões utilizadas do disco
            MarkUsedPartitionBlocks(
                this.CurrentPartition.AplOffset,
                this.CurrentPartition.AplSize);

            MarkUsedPartitionBlocks(
                this.CurrentPartition.DolOffset,
                this.CurrentPartition.FstOffset -
                this.CurrentPartition.DolOffset); //Tamanho do DOL = Posição do FST - Posição do DOL

            //Processa a tabela do sistema de arquivos
            this.CurrentPartition.Fst = new byte[this.CurrentPartition.FstSize];
            if ((e = ReadPartitionData(
                this.CurrentPartition.Fst,
                this.CurrentPartition.FstOffset,
                this.CurrentPartition.FstSize)) != WbfsError.Ok)
                goto RETURN;

            this.CurrentPartition.FstItemsCount =
                WbfsHelper.ntohi(this.CurrentPartition.Fst, 0x8);

            this.CurrentPartition.FstNameTablePos =
                0x0C * this.CurrentPartition.FstItemsCount; // 12 bytes de informação por arquivo

            //Processa os itens
            if (this.CurrentPartition.FstItemsCount > 1) // o item 0 descreve a Fst
                ParseFileSystem();

            //Limpa as informações desnecessárias
            this.aesDecryptor.Clear();
            this.aesDecryptor = null;
            this.CurrentPartition.Ticket = null;
            this.CurrentPartition.TMD = null;
            this.CurrentPartition.Cert = null;
            this.CurrentPartition.AplHeader = null;
            this.CurrentPartition.Fst = null;

        RETURN:
            return e;
        }


        //---------------------------------------------------------------------------------------------------
        // Decripts the title key
        //---------------------------------------------------------------------------------------------------
        private WbfsError ParseTitleKey()
        {
            WbfsError e = WbfsError.Ok;

            //Copia a chave a partir do Ticket
            byte[] titlekey = new byte[0x10];
            Array.Copy(this.CurrentPartition.Ticket, 0x01bf, titlekey, 0, 0x10);

            //Copia o vetor de inicializaçao a partir do Ticket
            byte[] iv = new byte[0x10];
            Array.Copy(this.CurrentPartition.Ticket, 0x1dc, iv, 0, 8);

            //Aloca o vetor da chave
            this.CurrentPartition.TitleKey = new byte[0x10];

            try
            {
                //Cria e configura o decriptador
                this.aesDecryptor = RijndaelManaged.Create();
                this.aesDecryptor.BlockSize = 128;
                this.aesDecryptor.Mode = CipherMode.CBC;
                this.aesDecryptor.Padding = PaddingMode.Zeros;
                this.aesDecryptor.Key = WbfsConstants.DiscCommonKey;
                this.aesDecryptor.IV = iv;

                //Decriptografa a chave
                ICryptoTransform aesDecryptor = this.aesDecryptor.CreateDecryptor();
                aesDecryptor.TransformBlock(titlekey, 0, 16, this.CurrentPartition.TitleKey, 0);
                aesDecryptor = null;
            }
            catch
            {
                e = WbfsError.DiscTitleDecryptFail;
                goto RETURN;
            }

        RETURN:

            titlekey = null;
            iv = null;

            return e;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads the Disc's file system
        //---------------------------------------------------------------------------------------------------
        private void ParseFileSystem()
        {
            byte[] fst = this.CurrentPartition.Fst;

            //Estrutura da FST:
            //
            // 0x0000 -          Informação sobre os arquivos (12 * numFiles)
            // (12 * numFiles) - Tabela com o nome de cada arquivo

            //Estrutura da informação de cada indice da FST (12 bytes):
            //
            // 0x00 - Posição do nome na tabela de nomes da fst - 4 bytes com 0x00ffffff de máscara (3 bytes efetivos),
            //        O primeiro byte excluído pela máscara é um sinalizador de arquivo usado
            // 0x04 - Posição do arquivo na partição - 4 bytes com bits rotacionados 2 vezes para a esquerda
            // 0x08 - Tamanho do arquivo na partição - 4 bytes

            uint fileSize = 0;
            long fileOffset = 0;

            //Por definição, o indice '0' contém informação da FST e não de um arquivo,
            //os dados do arquivo começam no indice '1'
            for (int i = 1; i < this.CurrentPartition.FstItemsCount; i++)
            {
                //Byte de uso do item
                if (fst[12 * i] != 0)
                    continue;

                //Pega a posição e o tamanho do arquivo na partição
                fileOffset = (long)WbfsHelper.ntohi(fst, 0x0C * i + 0x04) << 2;
                fileSize = (uint)WbfsHelper.ntohi(fst, 0x0C * i + 0x08);

                //Marca a região da partição correspondente ao arquivo
                //como usada
                MarkUsedPartitionBlocks(fileOffset, fileSize);
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Sets a block as used
        //---------------------------------------------------------------------------------------------------
        private void MarkUsedSectors(long offset, long length)
        {
            int block = (int)(offset >> WbfsConstants.WiiSectorSizeLSh);
            length += (offset % WbfsConstants.WiiSectorSize);

            while (length > 0)
            {
                this.WiiSectorUsage[block++] = true;
                length -= 0x8000;
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Sets a partition block as used
        //---------------------------------------------------------------------------------------------------
        private void MarkUsedPartitionBlocks(long offset, long length)
        {
            //Embora um bloco de dados ocupe 0x8000 bytes o sistema trabalha
            //com um tamanho efetivo de 0x7C00 por causa do alinhamento da
            //criptografia
            int pblock = this.CurrentPartition.DataStartBlock;
            int block = (int)(offset / 0x7C00);

            while (length > 0)
            {
                this.WiiSectorUsage[pblock + block++] = true;
                length -= 0x7C00;
            }
        }


        //---------------------------------------------------------------------------------------------------
        // Finds the last used partition block
        //---------------------------------------------------------------------------------------------------
        private int GetLastUsedWiiSector()
        {
            for (int i = this.WiiSectorUsage.Length - 1; i >= 0; i--)
                if (this.WiiSectorUsage[i]) return i;

            return -1;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads a data block from the disc
        //---------------------------------------------------------------------------------------------------
        private WbfsError ReadDisc(Byte[] buffer, long offset, int length)
        {
            WbfsError e = WbfsError.Ok;

            //Seta a posição do stream para o offset
            this.reader.BaseStream.Position = offset;

            //Le os dados e retorna em caso de erro
            if (this.reader.BaseStream.Read(buffer, 0, length) != length)
                e = WbfsError.IOReadError;

            //Marca a região como usada
            MarkUsedSectors(offset, length);

            return e;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads a data block from a partition
        //---------------------------------------------------------------------------------------------------
        private WbfsError ReadPartition(Byte[] buffer, long offset, int length)
        {
            //Le os dados na posição relativa à partição atual
            return ReadDisc(buffer, offset + this.CurrentPartition.Info.Position, length);
        }


        //---------------------------------------------------------------------------------------------------
        // Reads and decrypts a 32768 data block from the disc
        //---------------------------------------------------------------------------------------------------
        private WbfsError ReadCiphedBlock(long offset)
        {
            WbfsError e = WbfsError.Ok;

            //Lê os dados da partição
            if ((e = ReadPartition(this.ReadBlock, offset, 
                WbfsConstants.WiiSectorSize)) != WbfsError.Ok)
                return e;

            //Copia o vetor de inicialização
            byte[] iv = new byte[16];
            Array.Copy(this.ReadBlock, 0x3D0, iv, 0, 16);

            //Copia os dados criptografados
            byte[] ciphed = new byte[0x7C00];
            Array.Copy(this.ReadBlock, 0x400, ciphed, 0, 0x7C00);

            try
            {
                //Seta a chave e o vetor de inicialização
                this.aesDecryptor.Key = this.CurrentPartition.TitleKey;
                this.aesDecryptor.IV = iv;

                //Cria o decifrador e decriptografa os dados
                ICryptoTransform aesTransform = this.aesDecryptor.CreateDecryptor();
                aesTransform.TransformBlock(ciphed, 0, 0x7C00, this.ReadBlock, 0);
            }
            catch
            {
                e = WbfsError.DiscBlockDecryptFail;
                goto RETURN;
            }

        RETURN:

            iv = null;
            ciphed = null;

            return e;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads and decrypts a 32768 data block from a partition
        //---------------------------------------------------------------------------------------------------
        private WbfsError ReadPartitionData(Byte[] buffer, long offset, int length)
        {
            WbfsError e = WbfsError.Ok;

            //Embora o setor de Wii tenha 0x8000 bytes, parte deles contém informação para decriptação
            //do restante do setor e alinhamento, apenas 0x7C00 são dados reais e todos os offsets e 
            //tamanhos são calculados adotando 0x7C00 como o tamanho efetivo do setor.

            int sector = (int)(offset / 0x7C00);
            int read = 0;
            int toread;

            while (length > 0)
            {
                toread = (int)Math.Min(length, 0x7C00 - (offset % 0x7C00));

                e = ReadCiphedBlock(this.CurrentPartition.DataOffset + WbfsConstants.WiiSectorSize * sector);
                if (e != WbfsError.Ok) break;

                Array.Copy(this.ReadBlock, (int)(offset % 0x7C00), buffer, read, toread);

                read += toread;
                length -= toread;
                offset += toread;
                sector++;
            }

            return e;
        }


        //---------------------------------------------------------------------------------------------------
        // Reads an uint and convert it from network to host
        //---------------------------------------------------------------------------------------------------
        private uint ReadUInt32()
        {
            return (uint)IPAddress.NetworkToHostOrder(reader.ReadInt32());
        }


        //---------------------------------------------------------------------------------------------------
        // Reads an int and convert it from network to host
        //---------------------------------------------------------------------------------------------------
        private int ReadInt32()
        {
            return IPAddress.NetworkToHostOrder(reader.ReadInt32());
        }
    }
}
